import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

import torch
from efficientnet_pytorch import EfficientNet
from torchvision import transforms
from PIL import Image
import numpy as np
import cv2
import gradio as gr

# === 模型加载（仅首次） ===
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
checkpoint_path = './result/checkpoint_epoch_7.pth'
class_names = ['cat', 'dog', 'other']
model_arch = 'efficientnet-b0'

def load_model():
    model = EfficientNet.from_pretrained(model_arch, num_classes=len(class_names))
    model.load_state_dict(torch.load(checkpoint_path, map_location=device), strict=False)
    model.to(device)
    model.eval()
    return model

model = load_model()

# === Grad-CAM 辅助函数 ===
def generate_gradcam(model, input_tensor, target_class=None):
    gradients = []
    activations = []

    def backward_hook(module, grad_in, grad_out):
        gradients.append(grad_out[0])

    def forward_hook(module, input, output):
        activations.append(output)

    target_layer = model._blocks[-1]._project_conv
    forward_handle = target_layer.register_forward_hook(forward_hook)
    backward_handle = target_layer.register_backward_hook(backward_hook)

    model.zero_grad()
    output = model(input_tensor)
    if target_class is None:
        target_class = output.argmax(dim=1).item()

    class_score = output[:, target_class]
    class_score.backward()

    # === 移除 hook，防止多次重复注册 ===
    forward_handle.remove()
    backward_handle.remove()

    gradients = gradients[0].cpu().data.numpy()[0]
    activations = activations[0].cpu().data.numpy()[0]

    weights = np.mean(gradients, axis=(1, 2))
    cam = np.zeros(activations.shape[1:], dtype=np.float32)

    for i, w in enumerate(weights):
        cam += w * activations[i]

    cam = np.maximum(cam, 0)
    cam = cv2.resize(cam, (input_tensor.shape[3], input_tensor.shape[2]))
    cam -= cam.min()
    cam /= cam.max()
    return cam


def show_cam_on_image(img: np.ndarray, mask: np.ndarray) -> np.ndarray:
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img)
    cam = cam / np.max(cam)
    return np.uint8(255 * cam)

# === 推理函数 ===
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

def predict(image):
    image_pil = image.convert("RGB")
    image_tensor = transform(image_pil).unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(image_tensor)
        _, predicted = torch.max(output, 1)
        predicted_class = class_names[predicted.item()]

    # Grad-CAM
    cam_mask = generate_gradcam(model, image_tensor, predicted.item())

    img_np = np.array(image_pil.resize((256, 256))).astype(np.float32) / 255.0
    cam_result = show_cam_on_image(img_np, cam_mask)

    return f"预测结果：{predicted_class}", Image.fromarray(cam_result)

# === Gradio UI ===
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="上传图片"),
    outputs=[gr.Text(label="分类结果"), gr.Image(label="Grad-CAM 可视化")],
    title="猫狗识别 Demo",
    description="上传一张猫/狗/其他图片，查看模型预测结果与 Grad-CAM 热力图"
)

if __name__ == "__main__":
    demo.launch()
